(* $Id: PFOLEnv.ml 204 2009-03-04 14:28:18Z yann.regisgianas $ *)

(* Pangolin, a functional programming language for correct programs.
   Copyright (C) 2007, Yann Régis-Gianas, François Pottier.
   Contact: yann.regisgianas@gmail.com

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>. *)

open Misc
open PFOLSyntax
module X = XCoreSyntax
open XCoreSyntaxExt

(* An identifier in the source is mapped to a PFOL identifier and a PFOL
   type scheme. *)
type pfol_value_info =
    Position.t * value_identifier * type_scheme

(* A type variable is mapped to a PFOL type constructor and its arity. *)
type pfol_type_info =
    Position.t * type_identifier * int

(* A predicate is mapped to a PFOL identifier and
   its encoded declaration in PFOL. *)
type pfol_predicate_info =
    Position.t * predicate_identifier * predicate_declaration

(* A predicat may be associated to an identifier that represents it as
   an object (this is used by the HOL to FOL encoding). *)
type pfol_predicate_object_info =
    value_identifier

type env =
    {
      values	  : (X.id * pfol_value_info) list;
      types	  : (X.tid * pfol_type_info) list;
      predicates  : (X.pid * pfol_predicate_info) list;
      predobjects : (X.pid * pfol_predicate_object_info) list;
      facts	  : (string option * bool * predicate) list;
      idmap	  : var Identifier.Map.t;
      atomidmap   : Var.AtomIdMap.t
    }

let empty =
  {
    values      = [];
    types       = [];
    predicates  = [];
    predobjects = [];
    facts       = [];
    atomidmap   = Var.AtomIdMap.empty;
    idmap       = Identifier.Map.empty;
  }

let bind_atom x env =
  let aim = Var.AtomIdMap.add x env.atomidmap in
  let id = Var.AtomIdMap.find x aim in
    { env with
	idmap     = Identifier.Map.add id x env.idmap;
	atomidmap = aim
    }

(* Types. *)

let introduce_rigid_type_variable pos env (v, k) =
  let a = Var.Atom.freshb (atom_id v) in
    bind_atom a
      { env with types = (v, (pos, a, arity_of_kind k)) :: env.types }

let introduce_rigid_type_variables pos env vs =
  List.fold_left (introduce_rigid_type_variable pos) env vs

let remove_rigid_type_variables env vs =
  { env with
      types =
      List.fold_left (
	fun env v -> Misc.list_remove_assoc_if (X.Var.Atom.equal v) env)
	env.types vs
  }

let lookup_type_variable env v =
  try
    let (_, (_, a, _)) =
      List.find (fun (x, _) -> X.Var.Atom.equal x v) env.types
    in
      a
  with Not_found ->
    Printf.eprintf "internal error: unbound %s\n" (atom_as_string v);
    assert false

let type_parameters env vs =
  List.map (lookup_type_variable env) vs

(* Values. *)

let bind_value pos env x s =
  let a = Var.Atom.freshb (atom_id x) in
    (bind_atom a { env with values = (x, (pos, a, s)) :: env.values }, a)

let lookup env x =
  let (_, (_, a, _)) =
    List.find (fun (y, _) -> X.Var.Atom.equal y x) env.values
  in
    a

let type_scheme_of_id env x =
  let (_, (_, _, ty)) =
    List.find (fun (y, _) -> X.Var.Atom.equal y x) env.values
  in
    ty

let type_of_id env x =
  let (vs, ty) = open_type_scheme_abs (type_scheme_of_id env x) in
    assert (vs = []);
    ty

let remove_value x env =
  { env with
      values = List.filter (fun (_, (_, a, _)) ->
			      not (Var.Atom.equal a x))  env.values
  }

let remove_locals locals env =
  Var.AtomSet.fold remove_value locals env

(* Predicates. *)

let object_prefix id =
  let pos = PIdentifier.position id in
    PIdentifier.mk_value_id (pos, "object_" ^ PIdentifier.as_string id)

let wrap_object_prefix x =
  XCoreSyntax.Var.Atom.freshb (object_prefix (atom_id x))

let bind_predicate_object pos env x b =
  { env with predobjects = (x, b) :: env.predobjects }

let bind_value pos env x s =
  let a = Var.Atom.freshb (atom_id x) in
    (bind_atom a { env with values = (x, (pos, a, s)) :: env.values }, a)

let bind_predicate pos env x s d =
  let (env, b) = bind_value pos env (wrap_object_prefix x) s in
  let env = bind_predicate_object pos env x b in
  let a = Var.Atom.freshb (atom_id x) in
    (bind_atom a
       { env with predicates = (x, (pos, a, d)) :: env.predicates }, b, a)

let lookup_predicate_as_object env x =
  List.assoc x env.predobjects

let lookup_predicate env x =
  let (_, (_, a, _)) =
    List.find (fun (y, _) -> X.Var.Atom.equal y x) env.predicates
  in
    a

let is_predicate env x =
  List.exists (fun (y, _) -> X.Var.Atom.equal y x) env.predicates

let is_predicate_object env x =
  try
    let (_, (_, _, s)) =
      List.find (fun (y, _) -> X.Var.Atom.equal y x) env.values
    in
    let (_, ty) = open_type_scheme_abs s in
      match ty with
	| TPredicate _ -> true
	| TPArrow _ -> true
	| _ -> false
  with Not_found -> false

(* Facts. *)

let bind_fact ?mark pos status env f =
  { env with facts = (mark, status, f) :: env.facts }

let extract_type_parameters = function
  | ForallTys abs ->
      let (ts, t) = open_polyf_abs abs in
	(ts, t)
  | t ->
      ([], t)

let rec cons_axiom_declaration ?filter q (mark, status, g) =
  match filter with
    | Some filter when filter mark -> q
    | _ ->
        let (ts, f) = extract_type_parameters g in
        let q = ConsDecl (create_cons_decl_abs (DAxiom (ts, f), q)) in
          if status then q else ConsGoal (g, q)

let rec cons_value_declaration q (_, (_, v, s)) =
  ConsDecl (create_cons_decl_abs (DVal (v, s), q))

let rec cons_predicate_declaration q (_, (_, t, d)) =
  ConsDecl (create_cons_decl_abs (DPred (t, d), q))

let rec cons_type_declaration q (_, (_, t, k)) =
  ConsDecl (create_cons_decl_abs (DType (t, k), q))

let add_env_to_query ?filter env query =
  query
  |$ list_fdl_end (cons_axiom_declaration ?filter) env.facts
  |$ list_fdl_end cons_predicate_declaration env.predicates
  |$ list_fdl_end cons_value_declaration env.values
  |$ list_fdl_end cons_type_declaration env.types

let export_query query =
  try
    (export_query Var.AtomIdMap.empty query,
     export_query Var.AtomIdMap.empty (PFOLtoFOL.query query))
  with Var.Atom.Unknown a ->
    Printf.eprintf "Internal error: %s%d\n"
      (PIdentifier.as_string (Var.Atom.basename a))
      (Var.Atom.identity a);
    assert false

let env_to_query env =
  export_query (add_env_to_query env EmptyQuery)

let proof_obligation_to_query ?filter env f =
  export_query (add_env_to_query ?filter env (ConsGoal (f, EmptyQuery)))
